﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Management;
using System.Diagnostics;

namespace Confessor
{
    class Interrogation
    {
        xFileActions xFA = new xFileActions();
        List<ScanErrorObject> scanErrorsList = new List<ScanErrorObject>();
        /* Overall order in which this section resolves.
         * Check WMI connectivity, set the target host name.
         * Map the network drive.
         * Copy over the file.
         * Run the Confessor batch file on the remote machine.
         * Check to see if process is done (loop)
         * copy back results
         * - Cleanup
         * Delete data on remote machine
         * Ensure all processes have ended
         */

        //public ScanObject InitiateInterrogation(ScanObject scanObject, int y)
        public List<ScanErrorObject> InitiateInterrogation(ScanObject scanObject, int y)
        {
            //ScanObject tempScanObject = scanObject;
            ScanErrorObject scanError = new ScanErrorObject();
            // Set up class specific variables so confessor can proceed forward with each step.
            string targetHost = scanObject.Hosts[y].IP;
            string hostName = string.Empty;
            string mappedDriveLetter = string.Empty;
            bool wmiConnectivity = false;
            bool createdResultDirectory = false;
            bool mappedDrive = false;

            bool createdRemoteDirectoryStructure = false;
            bool copiedFilesToTarget = false;
            bool confessionStarted = false;
            bool createdTestProcess = false;
            bool finishedConfession = false;
            bool retrievedConfession = false;
            string resultsDirectory = string.Empty;
            
            // Test WMI Connectivty
            if (scanObject.WMIShares || scanObject.WMIFileCopy)
            {
                wmiConnectivity = testWMI(targetHost, wmiConnectivity, scanObject);
            }

            // Set the Hostname
            if (wmiConnectivity)
            {
                // Set the friendly name
                hostName = setTargetHostFriendlyName(targetHost, y, scanObject);
                scanObject.Hosts[y].HostName = hostName;
            }
            // setting it to targethost for PSExec where a friendly name isn't set.
            else
            {
                hostName = targetHost;
                scanObject.Hosts[y].HostName = hostName;
            }

            if (wmiConnectivity || scanObject.PsExec || scanObject.PsExecAsSystem)
            {
                // Set up the results directory
                createdResultDirectory = createTargetHostResultsDirectory(scanObject, y);
            }
            
            // The following mapping drive letter section does apply to WMIShares which relies on WMI exclusively.
            if ((createdResultDirectory && scanObject.WMIFileCopy) || (createdResultDirectory && scanObject.PsExec) || (createdResultDirectory && scanObject.PsExecAsSystem))
            {
                    // Map drive to target host
                mappedDriveLetter = mapNetworkDriveToTarget(scanObject, targetHost);
                    if (mappedDriveLetter.Length == 1)
                    {
                        mappedDrive = true;
                    }
                    // handle errors from call, if they exist.
                    else
                    {
                        xFA.WriteToErrorLog(hostName + ": Failed to Map Drive with error: " + mappedDriveLetter, scanObject);
                        addError(hostName, "Failed to Map Drive", mappedDriveLetter);
                        mappedDrive = false;
                    }

                    if (mappedDrive)
                    {
                        // Create remote directory structure
                        createdRemoteDirectoryStructure = createRemoteDirectories(mappedDriveLetter, hostName, scanObject);

                        if (createdRemoteDirectoryStructure)
                        {
                            // Copy the confessor files to target host.
                            copiedFilesToTarget = copyConfessorFilesToTarget(mappedDriveLetter, hostName, scanObject);
                        }
                    }
            }

            // Start the confession on target host
            if (copiedFilesToTarget || scanObject.WMIShares)
            {
                if (scanObject.WMITestCreateProcess)
                {
                    createdTestProcess = createTestProcess(targetHost, hostName, scanObject, y);
                }
                if (scanObject.WMIShares)
                {
                    confessionStarted = interrogateWMIShares(targetHost, hostName, scanObject, y);
                }
                if (scanObject.WMIFileCopy)
                {
                    confessionStarted = interrogateWMIMappedDrives(targetHost, hostName, y, scanObject);
                }
                if (scanObject.PsExec)
                {
                    confessionStarted = interrogatePSExec(targetHost, scanObject);
                }
                if (scanObject.PsExecAsSystem)
                {
                    confessionStarted = interrogatePSExecAsSystem(targetHost, scanObject);
                }

            }

            if (confessionStarted)
            {
                if (scanObject.WMIFileCopy || scanObject.PsExec || scanObject.PsExecAsSystem)
                {
                    // Wait for finish
                    finishedConfession = waitForConfessionToFinish(mappedDriveLetter, scanObject, scanObject.AllowedMinutesPerHost, y);

                    if (finishedConfession)
                    {
                        retrievedConfession = retrieveCompletedConfession(hostName, mappedDriveLetter, scanObject);
                    }
                }
                if (confessionStarted && scanObject.WMIShares)
                {
                    retrievedConfession = wmiSharesConfessionFinishedCheck(scanObject, scanObject.AllowedMinutesPerHost, y);
                }

            }
                 

            // Cleanup Mess
            cleanupMappedDrives(targetHost, hostName, mappedDriveLetter, scanObject, y, wmiConnectivity, 
                createdResultDirectory, mappedDrive, createdRemoteDirectoryStructure, 
                copiedFilesToTarget, confessionStarted, finishedConfession, retrievedConfession);

            return scanErrorsList;
        }

        // Test WMI Connectivity
        private bool testWMI(string targetHost, bool wmiConnectivity, ScanObject scanObject)
        {
            string result = xWMIActions.verifyWMIConnectivity(targetHost, scanObject.UserName, scanObject.Password);
            if (result == "true")
            {
                wmiConnectivity = true;
            }
            else
            {
                xFA.WriteToErrorLog(targetHost + ": Testing WMI Connectivity: " + result, scanObject);
                addError(targetHost, "Testing WMI Connectivity", result);
            }
            return wmiConnectivity;
        }

        // Set target host friendly name
        private string setTargetHostFriendlyName(string targetHost, int y, ScanObject scanObject)
        {
            string hostName = xWMIActions.returnTargetHostname(targetHost, scanObject.UserName, scanObject.Password);
            return hostName;
        }

        // Create Results directory for the target host
        private bool createTargetHostResultsDirectory(ScanObject scanObject, int y)
        {
            bool createdResultDirectory = false;
            // Set up the results directory
            try
            {
                //create directory for results of host scan
                string directory = scanObject.ResultsDirectory + @"\" + scanObject.Hosts[y].HostName;
                xFA.MakeDirectory(directory);
                createdResultDirectory = true;
            }
            catch (Exception ex)
            {
                xFA.WriteToErrorLog(scanObject.Hosts[y].HostName + ": Failed to create Results Directory : " + ex.Message, scanObject);
                addError(scanObject.Hosts[y].HostName, "Failed to create Results Directory", ex.Message);
                createdResultDirectory = false;
            }
            return createdResultDirectory;
        }

        // Map a network drive to the remote target system.
        private string mapNetworkDriveToTarget(ScanObject scanObject, string targetHost)
        {
            // Map the network drive to the target host
            string mappedDriveLetter = xNetworkDriveFunctions.mapNetworkDriveToTarget(scanObject.UserName, scanObject.Password, targetHost);

            return mappedDriveLetter;
        }

        // Create the directory structure on remote target system.
        private bool createRemoteDirectories(string mappedDriveLetter, string hostName, ScanObject scanObject)
        {
            bool createdRemoteDirectoryStructure = false;
            // Create Directory on remote host
            try
            {
                if (!Directory.Exists(mappedDriveLetter + @":\CONFESSOR"))
                {
                    Directory.CreateDirectory(mappedDriveLetter + @":\CONFESSOR");
                }
                if (!Directory.Exists(mappedDriveLetter + @":\CONFESSOR\results"))
                {
                    Directory.CreateDirectory(mappedDriveLetter + @":\CONFESSOR\results");
                }
                createdRemoteDirectoryStructure = true;
            }
            catch (Exception ex)
            {
                xFA.WriteToErrorLog(hostName + ": Failed to create directory structure on target host: " + ex.Message, scanObject);
                addError(hostName, "Failed to create directory structure on target host", ex.Message);
                createdRemoteDirectoryStructure = false;
            }
            return createdRemoteDirectoryStructure;
        }

        // Copy confessor and tools to remote target (For WMIFileCopy, PSExec, and PSExecAsSystem)
        private bool copyConfessorFilesToTarget(string mappedDriveLetter, string hostName, ScanObject scanObject)
        {
            bool copiedFilesToTarget = false;
            string status = xFA.CopyFile(scanObject.ResultsDirectory + @"\confess.bat", mappedDriveLetter + @":\CONFESSOR\confess.bat", true);
            if (status.Equals("success"))
            {
                copiedFilesToTarget = true;
            }
            else
            {
                xFA.WriteToErrorLog(hostName + ": Failed to copy confess.bat to target: " + status, scanObject);
                addError(hostName, "Failed to copy confess.bat to target", status);
                copiedFilesToTarget = false;
                return copiedFilesToTarget;
            }
            status = xFA.CopyFile(scanObject.ConfessorSharePath + @"\tools.zip", mappedDriveLetter + @":\CONFESSOR\tools.zip", true);
            if (status.Equals("success"))
            {
                copiedFilesToTarget = true;
            }
            else
            {
                xFA.WriteToErrorLog(hostName + ": Failed to copy tools.zip to target: " + status, scanObject);
                addError(hostName, "Failed to copy tools.zip to target", status);
                copiedFilesToTarget = false;
                return copiedFilesToTarget;
            }
            status = xFA.CopyFile(scanObject.ConfessorSharePath + @"\7za.exe", mappedDriveLetter + @":\CONFESSOR\7za.exe", true);
            if (status.Equals("success"))
            {
                copiedFilesToTarget = true;
            }
            else
            {
                xFA.WriteToErrorLog(hostName + ": Failed to copy 7za.exe to target: " + status, scanObject);
                addError(hostName, "Failed to copy 7za.exe to target", status);
                copiedFilesToTarget = false;
                return copiedFilesToTarget;
            }

            return copiedFilesToTarget;
        }

        // Kick off the confession of target host: WMI Shares
        private bool createTestProcess(string targetHost, string hostName, ScanObject scanObject, int y)
        {
            //
            bool confessionStarted = false;
            string commandstr = scanObject.WMICreateProcess;
            // Kick off the confession of target host
            string status = xWMIActions.startRemoteProcess(targetHost, scanObject.UserName, scanObject.Password, commandstr);
            if (status.Equals("success"))
            {
                scanObject.Hosts[y].ScanStartTime = DateTime.Now;
                confessionStarted = true;
            }
            else
            {
                xFA.WriteToErrorLog(hostName + ": Failed to start confession: " + status, scanObject);
                addError(hostName, "Failed to start confession", status);
                confessionStarted = false;
            }
            return confessionStarted;
        }

        // Kick off the confession of target host: WMI Shares
        private bool interrogateWMIShares(string targetHost, string hostName, ScanObject scanObject, int y)
        {
            //
            bool confessionStarted = false;
            string commandstr = scanObject.ResultsSharePath + @"\" + scanObject.InvestigationID + @"\confess.bat " + scanObject.ConfessorSharePath + " " + scanObject.ResultsSharePath + @"\" + scanObject.InvestigationID + @"\" + scanObject.Hosts[y].HostName;
            // Kick off the confession of target host
            string status = xWMIActions.startRemoteProcess(targetHost, scanObject.UserName, scanObject.Password, commandstr);
            if (status.Equals("success"))
            {
                scanObject.Hosts[y].ScanStartTime = DateTime.Now;
                confessionStarted = true;
            }
            else
            {
                xFA.WriteToErrorLog(hostName + ": Failed to start confession: " + status, scanObject);
                addError(hostName, "Failed to start confession", status);
                confessionStarted = false;
            }
            return confessionStarted;
        }

        // Kick off the confession of target host: WMI Mapped Drives
        private bool interrogateWMIMappedDrives(string targetHost, string hostName, int y, ScanObject scanObject)
        {
            bool confessionStarted = false;
            // Kick off the confession of target host
            string status = xWMIActions.startRemoteProcess(targetHost, scanObject.UserName, scanObject.Password, @"c:\confessor\confess.bat");
            if (status.Equals("success"))
            {
                scanObject.Hosts[y].ScanStartTime = DateTime.Now;
                confessionStarted = true;
            }
            else
            {
                xFA.WriteToErrorLog(hostName + ": Failed to start remote process: " + status, scanObject);
                addError(hostName, "Failed to start remote process", status);
                confessionStarted = false;
            }
            return confessionStarted;
        }

        // Kick off the confession of target host: PSExec
        private bool interrogatePSExec(string targetHost, ScanObject scanObject)
        {
            try
            {

                ProcessStartInfo p = new ProcessStartInfo(scanObject.ConfessorSharePath + @"\psexec.exe", @" \\" + targetHost + @" -u " + scanObject.UserName + @" -p " + scanObject.Password + @" c:\confessor\confess.bat");
                Process copyFilesProc = new Process();
                copyFilesProc.StartInfo = p;
                copyFilesProc.Start();
                copyFilesProc.WaitForExit();

                return true;
            }
            catch (Exception ex)
            {
                string statusMessage = ex.Message;
                xFA.WriteToErrorLog(targetHost + ": Failed to start PSExec: " + ex.Message, scanObject);
                addError(targetHost, "Failed to start PSExec", ex.Message);
                return false;
            }
        }

        // Kick off the confession of target host: PSExecAsSystem
        private bool interrogatePSExecAsSystem(string targetHost, ScanObject scanObject)
        {
            try
            {
                ProcessStartInfo p = new ProcessStartInfo(scanObject.ConfessorSharePath + @"\psexec.exe", @" \\" + targetHost + @" -s -u " + scanObject.UserName + @" -p " + scanObject.Password + @" c:\confessor\confess.bat");
                Process copyFilesProc = new Process();
                copyFilesProc.StartInfo = p;
                copyFilesProc.Start();
                copyFilesProc.WaitForExit();

                return true;
            }
            catch (Exception ex)
            {
                xFA.WriteToErrorLog(targetHost + ": Failed to start PSExec as system: " + ex.Message, scanObject);
                addError(targetHost, "Failed to start PSExec as system", ex.Message);
                return false;
            }
        }

        // Wait for scan to finish... loop here for x minutes.
        private bool waitForConfessionToFinish(string mappedDriveLetter, ScanObject scanObject, int minutes, int y)
        {
            bool finishedConfession = false;
            // Wait for finish
            int x = 0;
            int seconds = minutes * 60;
            while (x <= seconds)
            {
                System.Threading.Thread.Sleep(3000);

                if (File.Exists(mappedDriveLetter + @":\confessor\results\_finished.txt"))
                {
                    finishedConfession = true;
                    scanObject.Hosts[y].ScanFinishedTime = DateTime.Now;
                    scanObject.Hosts[y].ElapsedTime = scanObject.Hosts[y].ScanFinishedTime.Subtract(scanObject.Hosts[y].ScanStartTime);
                    break;
                }
                x++;
            }
            return finishedConfession;
        }

        private bool wmiSharesConfessionFinishedCheck(ScanObject scanObject, int seconds, int y)
        {
            bool finishedConfession = false;
            
            int x = 0;

            while (x <= seconds)
            {
                if (File.Exists(scanObject.ResultsSharePath + @"\" + scanObject.InvestigationID + @"\" + scanObject.Hosts[y].HostName + @"\_finished.txt"))
                {
                    scanObject.Hosts[y].ScanFinishedTime = DateTime.Now;
                    scanObject.Hosts[y].ElapsedTime = scanObject.Hosts[y].ScanFinishedTime.Subtract(scanObject.Hosts[y].ScanStartTime);
                    finishedConfession = true;
                    break;
                }
            }
            return finishedConfession;
        }

        // Retrieve confession
        private bool retrieveCompletedConfession(string hostName, string mappedDriveLetter, ScanObject scanObject)
        {
            bool retrievedConfession = false;
            //Copy Back Results
            string status = xFA.CopyFile(mappedDriveLetter + @":\CONFESSOR\results\results.zip", scanObject.ResultsDirectory + @"\" + hostName + @"\results.zip", true);
            if(status.Equals("success"))
            {
                retrievedConfession = true;
            }
            return retrievedConfession;
        }

        // Cleanup Mapped Drives
        private void cleanupMappedDrives(string targetHost, string hostName, string mappedDriveLetter, ScanObject scanObject, int y, bool wmiConnectivity,
            bool createdResultDirectory, bool mappedDrive, bool createdRemoteDirectoryStructure, bool copiedFilesToTarget,
            bool confessionStarted, bool finishedConfession, bool retrievedConfession)
        {
            string status = string.Empty;
            if (mappedDrive)
            {
                if (createdRemoteDirectoryStructure)
                {
                    // Remove Remote Confessor Files and Directory Structure
                    status = xFA.deleteDirectory(mappedDriveLetter + @":\CONFESSOR");
                    if (status.Equals("success"))
                    {
                    }
                    else
                    {
                        xFA.WriteToErrorLog(hostName + ": Cleanup - Could not delete all files and directory structure on remote system, manual Cleanup may be needed: " + status, scanObject);
                        addError(hostName, "Cleanup - Could not delete all files and directory structure on remote system", status);
                    }
                }

                // Remove mapped Network Drive
                status = xNetworkDriveFunctions.removeMappedDrive(mappedDriveLetter);
                if (status.Equals("success"))
                {
                }
                else
                {
                    xFA.WriteToErrorLog(hostName + ": Cleanup - Could not remove Mapped Network Drive" + mappedDriveLetter + ": " + status, scanObject);
                    addError(hostName, "Cleanup - Could nto remove Mapped Network Drive", mappedDriveLetter);
                }
            }
            if (retrievedConfession)
            {
                scanObject.Hosts[y].Status = "success";
                xFA.WriteToHostStatusLog(hostName + ": Completed successfully", scanObject);
            }
            else
            {
                scanObject.Hosts[y].Status = "failed";
                xFA.WriteToHostStatusLog(hostName + ": Failed to complete", scanObject);
            }
        }

        private void addError(string hostName, string confessorFailure, string error)
        {
            scanErrorsList.Add(new ScanErrorObject() { HostName = hostName, ConfessorFailure = confessorFailure, Error = error });
        }
    }
}
